package me.tzk.pdi.plugins.trans.steps.rocketmq;

import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.pentaho.di.trans.streaming.common.BlockingQueueStreamSource;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

import static java.nio.charset.Charset.defaultCharset;

/**
 * 注意：
 * kettle 底层框架使用的是 RxJava （Reactive 异步框架）
 * BlockingQueueStreamSource 中又背压机制，但是 DefaultMQPushConsumer 中的消费线程太多，需要调小
 */
public class RocketMQStreamSource extends BlockingQueueStreamSource<List<Object>> {
    private static final Class<?> PKG = RocketMQStreamSource.class;
    private final RocketMQConsumerMeta rocketMQConsumerMeta;
    private final RocketMQConsumer rocketMQConsumer;

    protected DefaultMQPushConsumer client;

    private final MessageListenerConcurrently callback = (messageList, consumeContext) -> {
        List collect = messageList.stream()
                .map(messageExt -> Arrays.asList(
                        messageExt.getMsgId(),
                        messageExt.getTopic(),
                        messageExt.getTags(),
                        new String(messageExt.getBody(), defaultCharset())
                )).collect(Collectors.toList());
        // 传递到下一个节点（这里是异步的，原因见上面）
        acceptRows(collect);
        return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
    };

    RocketMQStreamSource(RocketMQConsumerMeta rocketMQConsumerMeta, RocketMQConsumer rocketMQConsumer) {
        super(rocketMQConsumer);
        this.rocketMQConsumerMeta = rocketMQConsumerMeta;
        this.rocketMQConsumer = rocketMQConsumer;
    }

    @Override
    public void open() {
        try {
            client = RocketMQClientBuilder.build()
                    .withStep(rocketMQConsumer)
                    .withONSAddr(rocketMQConsumerMeta.getONSAddr())
                    .withCustomerId(rocketMQConsumerMeta.getConsumerId())
                    .withTopic(rocketMQConsumerMeta.getTopic())
                    .withTags(rocketMQConsumerMeta.getTags())
                    .withAccessKey(rocketMQConsumerMeta.getAccessKey())
                    .withSecretKey(rocketMQConsumerMeta.getSecretKey())
                    .withCallback(callback)
                    .withConsumeThreadNums(rocketMQConsumerMeta.getConsumeThreadNums())
                    .withMessageModel(rocketMQConsumerMeta.getMessageModel())
                    .withConsumeFromWhere(rocketMQConsumerMeta.getConsumeFromWhere())
                    .createConsumer();

            client.start();
            rocketMQConsumer.logBasic("RocketMQ stream source starts success.");
        } catch (Exception e) {
            rocketMQConsumer.stopAll();
            rocketMQConsumer.logError(e.toString());
        }
    }

    @Override
    public void close() {
        rocketMQConsumer.logBasic("RocketMQ stream source close!!!.");
        try {
            if (client != null) {
                client.shutdown();
            }
        } catch (Exception exception) {
            rocketMQConsumer.logError(exception.getMessage());
        } finally {
            super.close();
        }
    }

}
